use std::sync::Arc;

use datafusion::arrow::array::{Array, GenericStringArray, UInt64Array};
use datafusion::arrow::datatypes::DataType;
use datafusion::common::{Column, ScalarValue};
use datafusion::logical_expr::{ScalarUDF, Volatility};
use datafusion::logical_expr::expr::ScalarFunction;
use datafusion::physical_plan::ColumnarValue;
use datafusion::prelude::{col, create_udf, Expr};

use symbol::Symbol;

pub(super) fn symbol_to_string_utf() -> ScalarUDF {
    create_udf(
        "symbol_to_string",
        vec![DataType::UInt64],
        Arc::new(DataType::Utf8),
        Volatility::Stable,
        Arc::new(symbol_to_string),
    )
}

pub fn call_symbol_to_string(name: impl Into<Column>) -> Expr {
    Expr::ScalarFunction(ScalarFunction::new_udf(
        Arc::new(symbol_to_string_utf()),
        vec![col(name)],
    ))
}

fn symbol_to_string(args: &[ColumnarValue]) -> datafusion::error::Result<ColumnarValue> {
    assert_eq!(1, args.len(), "symbol_to_string args len must be 1");
    let arg = &args[0];
    match arg {
        ColumnarValue::Array(array) => {
            unsafe {
                compute_array(array.as_any().downcast_ref_unchecked::<UInt64Array>())
            }
        }
        ColumnarValue::Scalar(s) => {
            match s {
                ScalarValue::UInt64(None) => {
                    Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)))
                }
                ScalarValue::UInt64(Some(v)) => {
                    let tag_str = unsafe {
                        Symbol::from(*v).get_string_ref().to_string()
                    };
                    Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(tag_str))))
                }
                _ => {
                    Err(datafusion::error::DataFusionError::Execution(format!(
                        "symbol_to_string arg: {:?} not support", arg
                    )))
                }
            }
        }
    }
}

fn compute_array(a: &UInt64Array) -> datafusion::error::Result<ColumnarValue> {
    let str_array: GenericStringArray<i32> = a
        .into_iter()
        .map(|o| o.map(|v| {
            unsafe { Symbol::from(v).get_string_ref() }.to_string()
        }))
        .collect::<Vec<Option<String>>>()
        .into();
    Ok(ColumnarValue::Array(Arc::new(str_array)))
}

#[cfg(test)]
mod test {
    use datafusion::arrow::array::{Array, GenericStringArray, UInt64Array};
    use datafusion::logical_expr::ColumnarValue;

    use symbol::to_symbol;

    use crate::udf::symbol_to_string::compute_array;

    #[test]
    fn test() {
        let n = 10;
        let a0: UInt64Array = (0..n)
            .map(|i| {
                if i % 2 == 0 {
                    let s = format!("str_{}", i);
                    let v = to_symbol(&*s).unwrap().get_value();
                    Some(v)
                } else {
                    None
                }
            })
            .collect::<Vec<Option<u64>>>()
            .into();
        let a1 = compute_array(&a0).unwrap();
        let a2 = match &a1 {
            ColumnarValue::Array(array) => {
                unsafe { array.as_any().downcast_ref_unchecked::<GenericStringArray<i32>>() }
            }
            ColumnarValue::Scalar(_) => {
                unreachable!("a1 can not be a scalar value")
            }
        };
        for s in a2 {
            println!("s: {:?}", s)
        }
    }
}